package in.ac.iitb.cse.nlp.poswordpredictor;

import in.ac.iitb.cse.nlp.postagger.data.DataMaps;
import in.ac.iitb.cse.nlp.postagger.data.Tag;
import in.ac.iitb.cse.nlp.postagger.data.TransitionMatrix;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;

public class ExtendedTransitionMatrix extends TransitionMatrix {

	List<Prediction> getSuitablePrediction(String tagString) {
		Tag tag = DataMaps.tags.get(tagString);
		List<Prediction> predictions = new ArrayList<Prediction>();
		if (tag != null) {
			HashMap<Tag, Integer> hashMap = transitionTable.get(tag);
			List<TagPrediction> temp = new ArrayList<TagPrediction>();
			if (hashMap != null) {
				for (Tag key : hashMap.keySet()) {
					temp.add(new TagPrediction(key.getString(), getProbability(
							tag.getString(), key.getString())));
				}
			}
			Collections.sort(temp);
			for (int i = 0; i < temp.size()
					&& i < Constants.MAX_NUM_OF_PREDICTONS; i++) {
				predictions.add(new Prediction(temp.get(i).tag,
						temp.get(i).probability));
			}
		}
		return predictions;
	}

	@Override
	public String toString() {
		Tag tag = new Tag("^");
		StringBuffer buffer = new StringBuffer();
		for (Tag currentTag : transitionTable.get(tag).keySet()) {
			buffer.append("(" + "^" + "," + currentTag + ")="
					+ transitionTable.get(tag).get(currentTag) + ":"
					+ getProbability("^", currentTag.getString()));
			buffer.append("\n");
		}
		return buffer.toString();
	}

}

class TagPrediction implements Comparable<TagPrediction> {

	String tag;
	Double probability;

	TagPrediction(String tag, Double probability) {
		this.tag = tag;
		this.probability = probability;
	}

	@Override
	public int compareTo(TagPrediction prediction) {
		// We want to sort in descending order
		if (probability < prediction.probability) {
			return 1;
		} else if (probability > prediction.probability) {
			return -1;
		} else {
			return 0;
		}
	}
}